import random

import numpy as np
from sklearn import linear_model
import scipy.misc
from matplotlib import pyplot as plt
import MainCode.TOOL.ToolClass as tool
from skimage.metrics import structural_similarity as sk_cpt_ssim


class KSVD(object):
    def __init__(self, n_components, max_iter=50, tol=1e-6,
                 n_nonzero_coefs=None):
        """
        稀疏模型Y = DX，Y为样本矩阵，使用KSVD动态更新字典矩阵D和稀疏矩阵X
        :param n_components: 字典所含原子个数（字典的列数）
        :param max_iter: 最大迭代次数
        :param tol: 稀疏表示结果的容差
        :param n_nonzero_coefs: 稀疏度
        """
        self.dictionary = None
        self.sparsecode = None
        self.max_iter = max_iter
        self.tol = tol
        self.n_components = n_components
        self.n_nonzero_coefs = n_nonzero_coefs

    def _initialize(self, y):
        """
        初始化字典矩阵
        """
        u, s, v = np.linalg.svd(y)
        self.dictionary = u[:, :self.n_components]

    def _update_dict(self, y, d, x):
        """
        使用KSVD更新字典的过程
        """
        for i in range(self.n_components):
            index = np.nonzero(x[i, :])[0]
            if len(index) == 0:
                continue

            d[:, i] = 0
            r = (y - np.dot(d, x))[:, index]
            u, s, v = np.linalg.svd(r, full_matrices=False)
            d[:, i] = u[:, 0].T
            x[i, index] = s[0] * v[0, :]
        return d, x

    def fit(self, y):
        """
        KSVD迭代过程
        """
        self._initialize(y)
        for i in range(self.max_iter):
            x = linear_model.orthogonal_mp(self.dictionary, y, n_nonzero_coefs=self.n_nonzero_coefs)
            e = np.linalg.norm(y - np.dot(self.dictionary, x))
            if e < self.tol:
                break
            self._update_dict(y, self.dictionary, x)

        self.sparsecode = linear_model.orthogonal_mp(self.dictionary, y, n_nonzero_coefs=self.n_nonzero_coefs)
        return self.dictionary, self.sparsecode

def normalization(data):
    _range = np.max(abs(data))
    return data / _range

if __name__ == '__main__':
    font_size = 12
    # data_path = r'../Data/Numpy_DATA/AddNoise/theoretical_1NoiseData1.npy'
    # data_ori_path = r'../Data/Numpy_DATA/Original/theoreticalData1.npy'
    # np_data = np.load(data_path, allow_pickle=True)

    testdata_full = tool.__LoadData__(r'dataSource2.txt', 192, 800)
    testdata_full = normalization(testdata_full.transpose())
    plt.xlabel('Trace number', fontsize=font_size)
    plt.ylabel('Samples', fontsize=font_size)
    im2 = plt.imshow(testdata_full, aspect='auto', cmap='seismic', vmin=-1.0, vmax=1.0)
    plt.savefig('Real—original.pdf', format='pdf')
    plt.show()
    # testdata_full = np.load('vaild_data_01.npy', allow_pickle=True)
    # SNR = 10
    # noise = np.random.randn(testdata_full.shape[0], testdata_full.shape[1])  # 产生N(0,1)噪声数据
    # noise = noise - np.mean(noise)  # 均值为0
    # signal_power = np.linalg.norm(testdata_full - testdata_full.mean()) ** 2 / testdata_full.size  # 此处是信号的std**2
    # noise_variance = signal_power / np.power(10, (SNR / 10))  # 此处是噪声的std**2
    # noise = (np.sqrt(noise_variance) / np.std(noise)) * noise  ##此处是噪声的std**2
    # data_input = noise + testdata_full
    #
    #
    #
    # # np_ori_data = np.load(data_ori_path, allow_pickle=True)
    # # im_ascent = scipy.misc.ascent().astype(np.float)
    # PSNR1 = tool.snr(data_input, testdata_full)
    # ssim1 = sk_cpt_ssim(data_input, testdata_full)
    # print('去噪前, 信噪比为:{}'.format(PSNR1))
    # print('去噪前, 结构相似性：{}'.format(ssim1))
    # signal_noise = np.load('RealData2npy.npy', allow_pickle=True)
    # plt.xlabel('Trace number', fontsize=font_size)
    # plt.ylabel('Samples', fontsize=font_size)
    # im2 = plt.imshow(signal_noise, aspect='auto', cmap='seismic', vmin=-1.0, vmax=1.0)
    # plt.show()
    # plt.title('Original')
    # plt.xlabel('Trace number', fontsize=font_size)
    # plt.ylabel('Samples', fontsize=font_size)
    # im2 = plt.imshow(data_input, aspect='auto', cmap='seismic', vmin=-1.0, vmax=1.0)
    # cb2 = plt.colorbar(im2)
    # cb2.set_ticks([-1, -0.8, -0.6, -0.4, -0.2, 0, 0.2, 0.4, 0.6, 0.8, 1])
    # cb2.update_ticks()
    # plt.show()

    # for i in range(70,400,10):
    #     print()
    #     print(i)
    #     ksvd = KSVD(i)
    #     dictionary, sparsecode = ksvd.fit(signal_noise)
    #     # plt.figure()
    #     # plt.title('noisy data')
    #     # plt.xlabel('Trace number', fontsize=font_size)
    #     # plt.ylabel('Samples', fontsize=font_size)
    #     # im2 = plt.imshow(np_data, aspect='auto', cmap='seismic', vmin=-1.0, vmax=1.0)
    #     # cb2 = plt.colorbar(im2)
    #     # cb2.set_ticks([-1, -0.8, -0.6, -0.4, -0.2, 0, 0.2, 0.4, 0.6, 0.8, 1])
    #     # cb2.update_ticks()
    #     # plt.show()
    #     plt.xlabel('Trace number', fontsize=12)
    #     plt.ylabel('Samples', fontsize=12)
    #     im2 = plt.imshow(dictionary.dot(sparsecode), aspect='auto', cmap='seismic', vmin=-1.0, vmax=1.0)
    #     # cb2 = plt.colorbar(im2)
    #     # tick_locator = ticker.MaxNLocator(nbins=5)
    #     # cb2.locator = tick_locator
    #     # cb2.set_ticks([-1, -0.8, -0.6, -0.4, -0.2, 0, 0.2, 0.4, 0.6, 0.8, 1])
    #     # cb2.update_ticks()
    #     plt.savefig('no-{}-source-snr5-KSVD.pdf'.format(2), format='pdf')
    #     plt.show()
    #
    #     plt.xlabel('Trace number', fontsize=12)
    #     plt.ylabel('Samples', fontsize=12)
    #     im2 = plt.imshow(signal_noise - dictionary.dot(sparsecode), aspect='auto', cmap='seismic', vmin=-1.0, vmax=1.0)
    #     # cb2 = plt.colorbar(im2)
    #     # tick_locator = ticker.MaxNLocator(nbins=5)
    #     # cb2.locator = tick_locator
    #     # cb2.set_ticks([-1, -0.8, -0.6, -0.4, -0.2, 0, 0.2, 0.4, 0.6, 0.8, 1])
    #     # cb2.update_ticks()
    #     plt.savefig('no-{}-noise-snr5-KSVD.pdf'.format(2), format='pdf')
    #     plt.show()
        #
        #
        # plt.title('KSVD')
        # plt.xlabel('Trace number', fontsize=font_size)
        # plt.ylabel('Samples', fontsize=font_size)
        # im2 = plt.imshow(dictionary.dot(sparsecode), aspect='auto', cmap='seismic', vmin=-1.0, vmax=1.0)
        # cb2 = plt.colorbar(im2)
        # # cb2.set_ticks([-1, -0.8, -0.6, -0.4, -0.2, 0, 0.2, 0.4, 0.6, 0.8, 1])
        # # cb2.update_ticks()
        # plt.show()

        # PSNR1 = tool.__mtx_similar2__(np_data, testdata_full)
        # PSNR2 = tool.snr(dictionary.dot(sparsecode), testdata_full)
        # print('去噪后, 信噪比为:{}'.format(PSNR2))
        # ssim2 = sk_cpt_ssim(dictionary.dot(sparsecode), testdata_full)
        # print('去噪后, 结构相似性：{}'.format(ssim2))